#include "macroblock.h"
#include "h264_ctx.h"
#include "encoder_logger.h"
#include "analyse.h"
#include "dct.h"
#include "base.h"
#include <string.h>

enum
{
    PIXEL_16x16 = 0,
    PIXEL_16x8  = 1,
    PIXEL_8x16  = 2,
    PIXEL_8x8   = 3,
    PIXEL_8x4   = 4,
    PIXEL_4x8   = 5,
    PIXEL_4x4   = 6,

    /* Subsampled chroma only */
    PIXEL_4x16  = 7,  /* 4:2:2 */
    PIXEL_4x2   = 8,
    PIXEL_2x8   = 9,  /* 4:2:2 */
    PIXEL_2x4   = 10,
    PIXEL_2x2   = 11,
};

// yuv420亮度分量对应的chroma分量像素块大小映射表
static const uint8_t yuv420_luma2chroma_pixel[] = 
    { PIXEL_8x8,   PIXEL_8x4,  PIXEL_4x8,  PIXEL_4x4, PIXEL_4x2, PIXEL_2x4, PIXEL_2x2 };

// 扫描填充dct_dc
static const uint8_t block_idx_xy_1d[16] =
{
    0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15
};

static void macroblock_cache_load_neighbours(h264_ctx_t *ctx, int mb_x, int mb_y)
{
    macro_block_neighbors_info_t *neighbors = &ctx->cur_mb.neighbors;
    if (mb_x > 0) {
        neighbors->i_neighbour |= MB_LEFT;
    }
    int top_y = mb_y - 1;
    int top = top_y * ctx->cur_mb.i_mb_width + mb_x;
    if (top >= 0) {
        neighbors->i_neighbour |= MB_TOP;
    }

    if (mb_x > 0 && top_y >= 0) {
        neighbors->i_neighbour |= MB_TOPLEFT;
    }

    if(mb_x < ctx->cur_mb.i_mb_width - 1 && top_y >= 0) {
        neighbors->i_neighbour |= MB_TOPRIGHT;
    }
}

static void copy_pixel_to_enc_buff(uint8_t *dst, uint8_t *src, int i_stride, int i_width, int i_height)
{
    for (size_t i = 0; i < i_height; i++)
    {
        memcpy(dst, src, i_width);
        src += i_stride;
        dst += FENC_STRIDE;
    }
}

static void macroblock_load_pic_ptr(h264_ctx_t *ctx, int mb_x, int mb_y, int b_chroma)
{
    // copy luma
    if (b_chroma == 0) {
        pixel *p_src = ctx->encode_frame.plane_data[0] + 
            mb_x * H264_MB_SIZE + mb_y * ctx->cur_mb.i_mb_width * H264_MB_SIZE;
        copy_pixel_to_enc_buff((uint8_t *)ctx->cur_mb.pic.p_fenc[0], p_src
            , ctx->encode_frame.width, H264_MB_SIZE, H264_MB_SIZE);
        // TODO:copy fdec
    } else {
        // copy chroma
        pixel *p_src = ctx->encode_frame.plane_data[1] + 
            mb_x * H264_MB_SIZE / 2 + mb_y * ctx->cur_mb.i_mb_width * H264_MB_SIZE / 2;
        plane_copy_deinterleave(ctx->cur_mb.pic.p_fenc[1], FENC_STRIDE, 
                ctx->cur_mb.pic.p_fenc[1] +  FENC_STRIDE / 2, FENC_STRIDE, p_src,
                ctx->encode_frame.width, H264_MB_SIZE / 2, H264_MB_SIZE / 2);
        // TODO:copy fdec
    }
    
}

void macroblock_cache_load(h264_ctx_t *ctx, int mb_x, int mb_y)
{
    macroblock_cache_load_neighbours(ctx, mb_x, mb_y);
    INFO_LOG("macroblock_cache_load_neighbours: x:%d, y:%d neighbour %d\n", mb_x, mb_y, ctx->cur_mb.neighbors.i_neighbour);

    // first mb暂时不做
    macroblock_load_pic_ptr(ctx, mb_x, mb_y, 0);
    macroblock_load_pic_ptr(ctx, mb_x, mb_y, 1);

    // INFO_LOG("print enc buff:\n");
    // print_rect_buff_uint8((uint8_t *)ctx->cur_mb.pic.p_fenc[0], FENC_STRIDE, FENC_HEIGHT);
}


static void mb_analyse_intra(h264_ctx_t *ctx, h264_analyse_t *analyse)
{
    const intra16x16_pred_e * predict_mode = predict_16x16_mode_available(ctx->cur_mb.neighbors.i_neighbour);

    for ( ; *predict_mode != I_PRED_16x16_INVALID; predict_mode++ )
    {
        intra16x16_pred_e mode = *predict_mode;
        h264_predict_16x16(ctx->cur_mb.pic.p_fdec[0], mode);
        // print_rect_buff_uint8((uint8_t *)ctx->cur_mb.pic.p_fdec[0], FDEC_STRIDE, FDEC_HEIGHT - 2);
        int i_satd = x264_pixel_satd_16x16(ctx->cur_mb.pic.p_fenc[0], FENC_STRIDE, ctx->cur_mb.pic.p_fdec[0], FDEC_STRIDE);
        if( i_satd < analyse->i_satd_i16x16){
            analyse->i_satd_i16x16 = i_satd;
            analyse->i_intra16x16_pred = *predict_mode;
        }
    }
    
}

static void mb_analyse_intra_chroma(h264_ctx_t *ctx, h264_analyse_t *analyse)
{
    const intra_chroma_pred_e *predict_mode = predict_chroma_mode_available(ctx->cur_mb.neighbors.i_neighbour);
    int chromapix = yuv420_luma2chroma_pixel[PIXEL_16x16];
    int satd = 0;
    for( ; *predict_mode >= 0; predict_mode++ )
    {
        h264_predict_chroma(ctx->cur_mb.pic.p_fdec[1], *predict_mode);
        h264_predict_chroma(ctx->cur_mb.pic.p_fdec[2], *predict_mode);
        satd = x264_pixel_satd_8x8(ctx->cur_mb.pic.p_fenc[1], FENC_STRIDE, ctx->cur_mb.pic.p_fdec[1], FDEC_STRIDE);
        satd += x264_pixel_satd_8x8(ctx->cur_mb.pic.p_fenc[2], FENC_STRIDE, ctx->cur_mb.pic.p_fdec[2], FDEC_STRIDE);
        if( satd < analyse->i_satd_chroma){
            analyse->i_satd_chroma = satd;
            analyse->i_chroma_pred = *predict_mode;
        }
    }
}

static int x264_ctz_4bit( uint32_t x )
{
    static uint8_t lut[16] = {4,0,1,0,2,0,1,0,3,0,1,0,2,0,1,0};
    return lut[x];
}

#define FOREACH_BIT(idx,start,mask) for( int idx = start, msk = mask, skip; msk && (skip = x264_ctz_4bit(msk), idx += skip, msk >>= skip+1, 1); idx++ )


static void mb_encode_i16x16(h264_ctx_t *ctx, int i_qp)
{
    dctcoef dct4x4[16][16] = {0};
    dctcoef dct_dc4x4[16] = {0};
    pixel *p_src = ctx->cur_mb.pic.p_fenc[0];
    pixel *p_dst = ctx->cur_mb.pic.p_fdec[0];
    h264_predict_16x16(ctx->cur_mb.pic.p_fdec[0], ctx->cur_mb.i_intra16x16_pred_mode);

    // 16x16矩阵dct变换以4x4块为单位进行
    sub16x16_dct(dct4x4, p_src, p_dst);
    int nz = 0;
    int block_cbp = 0;
    memset(ctx->cur_mb.cache.non_zero_count, 0, sizeof(ctx->cur_mb.cache.non_zero_count));
    // print_rect_buff_uint8((uint8_t *)dct4x4, 16, 16, 16);

    // 提取dc系数
    for( int idx = 0; idx < 16; idx++ )
    {
        dct_dc4x4[block_idx_xy_1d[idx]] = dct4x4[idx][0];
        dct4x4[idx][0] = 0;
    }

    for( int i8x8 = 0; i8x8 < 4; i8x8++ )
    {
        // 一次性量化4个4x4块
        nz = quant_4x4x4(&dct4x4[i8x8*4], ctx->quant.quant4_mf[CQM_IY][i_qp], ctx->quant.quant4_bias[CQM_IY][i_qp] );
        if (nz != 0) {
            block_cbp = 0xf;
            int mask = nz;
            for( int idx = i8x8 * 4; mask != 0; idx++ )
            {
                int skip = x264_ctz_4bit(mask);
                idx += skip;
                mask >>= skip+1;
                zigzag_scan_4x4_frame(ctx->dct.luma4x4[idx], dct4x4[idx]);
                dequant_4x4(dct4x4[idx], ctx->quant.dequant4_mf[CQM_IY], i_qp);
                ctx->cur_mb.cache.non_zero_count[x264_scan8[idx]] = 1;
            }
        }
    }

    ctx->cur_mb.i_cbp_luma |= block_cbp;

    // 对dc系数进行dct变换+量化
    dct4x4dc(dct_dc4x4);
    nz = quant_4x4_dc(dct_dc4x4, ctx->quant.quant4_mf[CQM_IY][i_qp][0] >>1, 
            ctx->quant.quant4_bias[CQM_IY][i_qp][0] << 1);
    ctx->cur_mb.cache.non_zero_count[x264_scan8[LUMA_DC]] = nz;
    if (nz != 0) {
        // 拷贝亮度dc数据
        zigzag_scan_4x4_frame(ctx->dct.luma16x16_dc[0], dct_dc4x4);
        idct4x4dc(dct_dc4x4);
        dequant_4x4_dc(dct_dc4x4, ctx->quant.dequant4_mf[CQM_IY], i_qp);
        if( block_cbp != 0 )
            for( int i = 0; i < 16; i++ )
                dct4x4[i][0] = dct_dc4x4[block_idx_xy_1d[i]];
    }
    if (block_cbp != 0) {
        add16x16_idct(p_dst, dct4x4);
    } else if (nz != 0) {
        add16x16_idct_dc(p_dst, dct_dc4x4);
    }
    return;
}

void mb_encode_chroma( h264_ctx_t *ctx, int b_inter, int i_qp, int chroma422 )
{
    int (*dequant_mf)[16] = ctx->quant.dequant4_mf[CQM_IC + b_inter];
    dctcoef dct_dc[8] = {0};
    ctx->cur_mb.i_cbp_chroma = 0;
    int nz = 0;
    int nz_dc = 0;
    memset( ctx->cur_mb.cache.non_zero_count + 16, 0, 2);
    memset( ctx->cur_mb.cache.non_zero_count + 18, 0, 2);
    memset( ctx->cur_mb.cache.non_zero_count + 32, 0, 2);
    memset( ctx->cur_mb.cache.non_zero_count + 34, 0, 2);

    if (chroma422 == 1) {
        memset( ctx->cur_mb.cache.non_zero_count + 24, 0, 2);
        memset( ctx->cur_mb.cache.non_zero_count + 26, 0, 2);
        memset( ctx->cur_mb.cache.non_zero_count + 40, 0, 2);
        memset( ctx->cur_mb.cache.non_zero_count + 42, 0, 2);
    }

    for (size_t ch = 0; ch < 2; ch++) {
        pixel *p_src = ctx->cur_mb.pic.p_fenc[1 + ch];
        pixel *p_dst = ctx->cur_mb.pic.p_fdec[1 + ch];
        int nz_ac = 0;
        dctcoef dct4x4[8][16] = {0};
        for (size_t j = 0; j <= chroma422; j++) {
            sub8x8_dct(&dct4x4[4 * j], p_src + 8 * j * FENC_STRIDE, p_dst + 8 * j * FDEC_STRIDE );
        }

        if( chroma422 == 1) {
            dct2x4dc( dct_dc, dct4x4 );
        }
        else {
            dct2x2dc( dct_dc, dct4x4 );
        }
        
        /* calculate dct coeffs */
        for (size_t i8x8 = 0;i8x8 < (chroma422 ? 2 : 1); i8x8++) {
            nz = quant_4x4x4(&dct4x4[i8x8*4], ctx->quant.quant4_mf[CQM_IC + b_inter][i_qp],
                            ctx->quant.quant4_bias[CQM_IC + b_inter][i_qp]);
            nz_ac |= nz;
            int mask = nz;
            for( int i4x4 = 0,  skip; mask != 0; i4x4++ )
            {
                int skip = x264_ctz_4bit(mask);
                i4x4 += skip;
                mask >>= skip + 1;
                int idx = 16+ch*16+i8x8*8+i4x4;
                zigzag_scan_4x4_frame( ctx->dct.luma4x4[idx], dct4x4[i8x8*4+i4x4] );
                dequant_4x4( dct4x4[i8x8*4+i4x4], dequant_mf, i_qp );
                ctx->cur_mb.cache.non_zero_count[x264_scan8[idx]] = 1;
            }
        }

        nz_dc = 0;
        for( int i = 0; i <= chroma422; i++ ) {
            nz_dc |= quant_2x2_dc( &dct_dc[4*i], ctx->quant.quant4_mf[CQM_IC + b_inter][i_qp+3*chroma422][0] >> 1,
                                                 ctx->quant.quant4_bias[CQM_IC + b_inter][i_qp + 3*chroma422][0] << 1 );
        }

        ctx->cur_mb.cache.non_zero_count[x264_scan8[CHROMA_DC+ch]] = nz_dc;

        // 交流系数为0
        if (nz_ac == 0) {
            memset( ctx->cur_mb.cache.non_zero_count + 16 + 16 * ch, 0, 2);
            memset( ctx->cur_mb.cache.non_zero_count + 18 + 16 * ch, 0, 2);

            if (chroma422 == 1) {
                memset( ctx->cur_mb.cache.non_zero_count + 24 + 16 * ch, 0, 2);
                memset( ctx->cur_mb.cache.non_zero_count + 26 + 16 * ch, 0, 2);
            }

            // 量化后全为0
            if (nz_dc == 0) {
                continue;
            }

            // x264中有mb_optimize_chroma_dc，该步骤暂时省略
            if( chroma422 == 1 )
            {
                // zigzag_scan_2x4_dc( ctx->dct.chroma_dc[ch], dct_dc );
                // idct_dequant_2x4_dconly( dct_dc, dequant_mf, i_qp+3 );
            }
            else
            {
                zigzag_scan_2x2_dc( ctx->dct.chroma_dc[ch], dct_dc );
                idct_dequant_2x2_dconly( dct_dc, dequant_mf, i_qp );
            }

            for( int i = 0; i <= chroma422; i++ ) {
                add8x8_idct_dc( p_dst + 8*i*FDEC_STRIDE, &dct_dc[4*i] );
            }
        } else {
            ctx->cur_mb.i_cbp_chroma = 1;
            if( nz_dc != 0)
            {
                if( chroma422 == 1)
                {
                    // zigzag_scan_2x4_dc( h->dct.chroma_dc[ch], dct_dc );
                    // h->quantf.idct_dequant_2x4_dc( dct_dc, dct4x4, dequant_mf, i_qp+3 );
                }
                else
                {
                    zigzag_scan_2x2_dc( ctx->dct.chroma_dc[ch], dct_dc );
                    idct_dequant_2x2_dc( dct_dc, dct4x4, dequant_mf, i_qp );
                }
            }

            for( int i = 0; i <= chroma422; i++ ) {
                add8x8_idct( p_dst + 8*i*FDEC_STRIDE, &dct4x4[4*i] );
            }
        }
    }

    /* 0 = none, 1 = DC only, 2 = DC+AC */
    ctx->cur_mb.i_cbp_chroma += (ctx->cur_mb.cache.non_zero_count[x264_scan8[CHROMA_DC+0]] |
                           ctx->cur_mb.cache.non_zero_count[x264_scan8[CHROMA_DC+1]] | ctx->cur_mb.i_cbp_chroma);    
}

void macroblock_encode(h264_ctx_t *ctx)
{
    if( ctx->cur_mb.type == I_16x16 )
    {
        mb_encode_i16x16( ctx, ctx->cur_mb.i_qp);
    }

    // 当前程序必定编码色度块
    mb_encode_chroma( ctx, 1, ctx->cur_mb.i_qp, 0);
}

void marcoblock_analyse(h264_ctx_t *ctx)
{
    h264_analyse_t analyse = {0};
    mb_analyse_init(&analyse);
    if (ctx->slice_head.slice_type == SLICE_TYPE_I) {
        mb_analyse_intra(ctx, &analyse);
        // 目前只支持16x16
        ctx->cur_mb.type = I_16x16;
        ctx->cur_mb.i_intra16x16_pred_mode = analyse.i_intra16x16_pred;
    }
    mb_analyse_intra_chroma(ctx, &analyse);
    ctx->cur_mb.i_chroma_pred_mode = analyse.i_chroma_pred;
}